//Coefficient MatrixAB
#ifndef MATRIXAB_H_
#define MATRIXAB_H_
#include "./Step.h"
#include "./Constants.h"


class MatrixAB{
private:
    int m_nr, m_nt;
    double m_dz;
    double *A1B1_const;   //Used to store the same values ​​in arrays A1 and B1
    double *alfa_array;

    void set_alfa(double* frequency_array);
    void set_A1B1_const(double* r_grid, double* step_array);

    void set_A0B0(double* A0, double* B0, double* r_grid, double* step_array);
    void set_A1B1(complex** A1, complex** B1);
    void set_A2B2(double* A2, double* B2, double* step_array);

public: 
    explicit MatrixAB(int nr, int nt, double dz):m_nr(nr), m_nt(nt), m_dz(dz){
        A1B1_const  = new double[m_nr];
        alfa_array  = new double[m_nt];
    }

    virtual ~MatrixAB(){
        delete[] A1B1_const;
    }
    void operator()(double *frequency_array,
                    double *r_grid, double *step_array,
                    double *A0, double *B0,
                    complex **A1, complex **B1,
                    double *A2, double *B2);
    
  
};
void MatrixAB::set_alfa(double* frequency_array){
    double omega_i = 0;
    for(int i = 0; i < m_nt; i++){
        omega_i = 2. * M_PI * frequency_array[i];
        alfa_array[i] = 4. * omega_i/ (Constants::c);   
    }
}
void MatrixAB::set_A1B1_const(double* r_grid, double* step_array){
    
    for(int i = 0; i < m_nr; i++){
        double beta1 = 1. / (step_array[i] * step_array[i-1]);
        double beta2 = 1. / (r_grid[i] * step_array[i-1]);
        double beta3 = 1. / (step_array[i] * step_array[i]);
        beta1 *= m_dz;
        beta2 *= m_dz;
        beta3 *= m_dz;

        // A1B1_const[i] = (i == 0)? (2. * beta3) : (beta1 - beta2 + beta3);
        A1B1_const[i] = (i == 0)? (4 * beta3) : (beta1 - beta2 + beta3);//洛必达法则 
        
    }
}

void MatrixAB::set_A0B0(double* A0, double* B0, double* r_grid, double* step_array){
    for (int index = 0; index < m_nr; index++){
        double beta1 = 1. / (step_array[index] * step_array[index+1]);
        double beta2 = 1. / (r_grid[index+1] * step_array[index]);
        beta1 *= m_dz;
        beta2 *= m_dz;

        A0[index] = - beta1 +beta2;
        B0[index] = - A0[index];

    }     
}

void MatrixAB::set_A1B1(complex** A1, complex** B1)
{
    for(int i = 0; i < m_nt; i++)
    {  
        for(int j = 0;j < m_nr; j++)
        {
            A1[i][j] =   A1B1_const[j] + alfa_array[i] * Constants::Ii; 
            B1[i][j] = - A1B1_const[j] + alfa_array[i] * Constants::Ii;
        }
    }
}

void MatrixAB::set_A2B2(double* A2, double* B2, double* step_array){
    for(int index = 0; index < m_nr; index++){
        double beta3 = 1. / (step_array[index] * step_array[index]);
        beta3 *= m_dz;

        // A2[index]  = (index == 0)? (- 2. * beta3): (- beta3);
        A2[index]  = (index == 0)? (- 4. * beta3): (- beta3);
        B2[index]  = - A2[index];
    }
}



void MatrixAB::operator()(double *frequency_array,
                       double *r_grid, double *step_array,
                       double *A0, double *B0,
                       complex **A1, complex **B1,
                       double *A2, double *B2){
    set_alfa(frequency_array);
    set_A1B1_const(r_grid, step_array);
    set_A0B0(A0, B0, r_grid, step_array);
    set_A1B1(A1, B1);
    set_A2B2(A2, B2, step_array);
}

#endif